"""
改进版后端 - 修复文件上传问题
关键修复:
1. 添加文件上传接口
2. DirectStorySession 支持文件路径
3. 正确传递配置给 story_writer
"""

from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
from pathlib import Path
import os
import yaml
import json
import threading
import re
from datetime import datetime
import requests
import uuid
import pickle
from werkzeug.utils import secure_filename
import time
from threading import Lock

from mm_story_agent import MMStoryAgent
from mm_story_agent.base import init_tool_instance

app = Flask(__name__)
CORS(app)

# 文件上传配置
UPLOAD_FOLDER = Path("uploads")
UPLOAD_FOLDER.mkdir(exist_ok=True)
ALLOWED_EXTENSIONS = {'txt', 'pdf', 'docx', 'md', 'epub'}

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB

# 全局状态存储
user_sessions = {}

# 会话持久化目录
SESSION_STORE_PATH = Path("sessions")
SESSION_STORE_PATH.mkdir(exist_ok=True)

# 加载配置
try:
    with open('configs/mm_story_agent.yaml', 'r', encoding='utf-8') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    print("✅ 配置文件加载成功")
except Exception as e:
    print(f"❌ 配置文件加载失败: {e}")
    config = {}


def allowed_file(filename):
    """检查文件扩展名"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def save_session(session_id, session):
    """保存会话到文件"""
    try:
        with open(SESSION_STORE_PATH / f"{session_id}.pkl", 'wb') as f:
            pickle.dump(session, f)
        print(f"💾 会话已保存: {session_id}")
    except Exception as e:
        print(f"⚠️ 保存会话失败: {e}")


def load_session(session_id):
    """从文件加载会话"""
    try:
        session_file = SESSION_STORE_PATH / f"{session_id}.pkl"
        if session_file.exists():
            with open(session_file, 'rb') as f:
                session = pickle.load(f)
                print(f"📂 会话已恢复: {session_id}")
                return session
    except Exception as e:
        print(f"⚠️ 加载会话失败: {e}")
    return None


# class DashScopeClient:
#     """阿里云DashScope API客户端"""
#     def __init__(self, api_key: str):
#         self.api_key = api_key
#         self.base_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
        
#     def generate(self, prompt: str, system_prompt: str = None, **kwargs) -> str:
#         try:
#             headers = {
#                 "Authorization": f"Bearer {self.api_key}",
#                 "Content-Type": "application/json"
#             }
            
#             messages = []
#             if system_prompt:
#                 messages.append({"role": "system", "content": system_prompt})
#             messages.append({"role": "user", "content": prompt})
            
#             payload = {
#                 "model": kwargs.get("model", "qwen-turbo"),
#                 "input": {"messages": messages},
#                 "parameters": {
#                     "max_tokens": kwargs.get("max_tokens", 2000),
#                     "temperature": kwargs.get("temperature", 0.8),
#                     "top_p": 0.8
#                 }
#             }
            
#             response = requests.post(self.base_url, headers=headers, json=payload, timeout=60)
            
#             if response.status_code == 200:
#                 result = response.json()
#                 if result.get("output"):
#                     return result["output"]["text"]
#                 else:
#                     raise Exception(f"API返回格式错误: {result}")
#             else:
#                 raise Exception(f"API调用失败: {response.status_code}")
                
#         except Exception as e:
#             print(f"❌ DashScope调用错误: {str(e)}")
#             raise

class DashScopeClient:
    """阿里云DashScope API客户端 - 带速率限制"""
    
    # 类级别的锁和时间戳，用于全局速率限制
    _last_request_time = 0
    _request_lock = Lock()
    _min_interval = 0.6  # 🔥 关键：最小请求间隔（秒）- 从0增加到0.6
    
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.base_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
        self.retry_count = 3  # 🔥 重试次数
        self.retry_delay = 2.0  # 🔥 重试延迟（秒）
        
    def _wait_for_rate_limit(self):
        """等待以遵守速率限制"""
        with self._request_lock:
            current_time = time.time()
            time_since_last = current_time - self._last_request_time
            
            if time_since_last < self._min_interval:
                wait_time = self._min_interval - time_since_last
                print(f"⏳ 速率限制：等待 {wait_time:.2f} 秒...")
                time.sleep(wait_time)
            
            self._last_request_time = time.time()
    
    def generate(self, prompt: str, system_prompt: str = None, **kwargs) -> str:
        """生成文本 - 带重试和速率限制"""
        
        for attempt in range(self.retry_count):
            try:
                # 🔥 等待以遵守速率限制
                self._wait_for_rate_limit()
                
                headers = {
                    "Authorization": f"Bearer {self.api_key}",
                    "Content-Type": "application/json"
                }
                
                messages = []
                if system_prompt:
                    messages.append({"role": "system", "content": system_prompt})
                messages.append({"role": "user", "content": prompt})
                
                payload = {
                    "model": kwargs.get("model", "qwen-turbo"),
                    "input": {"messages": messages},
                    "parameters": {
                        "max_tokens": kwargs.get("max_tokens", 2000),
                        "temperature": kwargs.get("temperature", 0.8),
                        "top_p": 0.8
                    }
                }
                
                response = requests.post(
                    self.base_url, 
                    headers=headers, 
                    json=payload, 
                    timeout=60
                )
                
                if response.status_code == 200:
                    result = response.json()
                    if result.get("output"):
                        return result["output"]["text"]
                    else:
                        raise Exception(f"API返回格式错误: {result}")
                
                # 🔥 处理 429 速率限制错误
                elif response.status_code == 429:
                    result = response.json()
                    error_msg = result.get('message', 'Rate limit exceeded')
                    
                    if attempt < self.retry_count - 1:
                        wait_time = self.retry_delay * (attempt + 1)  # 指数退避
                        print(f"⚠️  速率限制 (尝试 {attempt + 1}/{self.retry_count})，等待 {wait_time} 秒后重试...")
                        time.sleep(wait_time)
                        continue
                    else:
                        raise Exception(f"速率限制，重试次数已用尽: {error_msg}")
                
                else:
                    raise Exception(f"API调用失败: HTTP {response.status_code}")
                    
            except Exception as e:
                if attempt < self.retry_count - 1:
                    print(f"⚠️  请求失败 (尝试 {attempt + 1}/{self.retry_count}): {str(e)}")
                    time.sleep(self.retry_delay)
                    continue
                else:
                    print(f"❌ DashScope调用错误: {str(e)}")
                    raise

class ConversationalStorySession:
    """对话式故事生成会话"""
    
    def __init__(self, session_id: str, api_key: str, target_pages: int = 8):
        self.session_id = session_id
        self.api_key = api_key
        self.target_pages = target_pages
        self.dialogue_history = []
        self.story_pages = []
        self.image_paths = []
        self.video_path = None
        self.current_stage = "init"
        self.is_generating = False
        self.progress = 0
        self.status_message = "准备就绪"
        self.llm = DashScopeClient(api_key)
        
        self.system_prompt = """你是一个专业的儿童故事创作助手。你的任务是:
1. 通过对话了解用户想要的故事类型、角色、场景等信息
2. 根据用户反馈不断优化故事构想
3. 当用户说"开始生成"或"可以了"时,总结信息并准备生成故事"""

    def update_progress(self, stage: str, progress: int, message: str):
        self.current_stage = stage
        self.progress = progress
        self.status_message = message
        print(f"📊 [{stage}] {progress}% - {message}")
    
    def chat(self, user_message: str, target_pages: int = None) -> dict:
        """处理用户消息"""
        try:
            if target_pages:
                self.target_pages = target_pages
            
            if any(keyword in user_message.lower() for keyword in 
                   ["开始生成", "生成故事", "可以了", "就这样", "开始吧"]):
                return self._prepare_generation()
            
            self.dialogue_history.append({
                "role": "user",
                "content": user_message,
                "timestamp": datetime.now().isoformat()
            })
            
            context = "\n".join([
                f"{msg['role']}: {msg['content']}" 
                for msg in self.dialogue_history[-6:]
            ])
            
            response = self.llm.generate(
                prompt=f"对话历史:\n{context}\n\n请继续对话:",
                system_prompt=self.system_prompt,
                temperature=0.8
            )
            
            self.dialogue_history.append({
                "role": "assistant",
                "content": response,
                "timestamp": datetime.now().isoformat()
            })
            
            self.current_stage = "chatting"
            
            return {
                "success": True,
                "response": response,
                "stage": self.current_stage,
                "can_generate": len(self.dialogue_history) >= 4
            }
            
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "response": "抱歉,我遇到了一些问题。请重试。"
            }
    
    def _prepare_generation(self) -> dict:
        self.current_stage = "preparing"
        story_params = self._extract_story_params()
        story_params["target_pages"] = self.target_pages
        
        return {
            "success": True,
            "response": f"好的!我现在根据我们的讨论来创作 {self.target_pages} 页的故事...",
            "stage": "preparing",
            "story_params": story_params
        }
    
    def _extract_story_params(self) -> dict:
        try:
            dialogue_text = "\n".join([
                f"{msg['role']}: {msg['content']}" 
                for msg in self.dialogue_history
            ])
            
            extract_prompt = f"""根据以下对话,提取故事的关键信息:

{dialogue_text}

请以JSON格式输出(不要包含markdown标记):
{{
    "story_topic": "故事的核心主题",
    "main_role": "主角的名字和特点",
    "scene": "故事发生的场景",
    "style": "故事风格"
}}

只输出JSON,不要有其他内容:"""
            
            result = self.llm.generate(prompt=extract_prompt, temperature=0.3)
            result = result.strip()
            result = re.sub(r'```json\s*', '', result)
            result = re.sub(r'```\s*', '', result)
            
            params = json.loads(result)
            params["input_type"] = "topic"
            return params
            
        except Exception as e:
            print(f"⚠️ 参数提取失败,使用默认值: {e}")
            return {
                "story_topic": "一个温暖有趣的冒险故事",
                "main_role": "勇敢的小主角",
                "scene": "奇幻世界",
                "input_type": "topic"
            }
    
    def generate_story(self, story_params: dict) -> list:
        """生成完整故事"""
        try:
            self.is_generating = True
            self.current_stage = "generating"
            self.progress = 30
            
            prompt = f"""请创作一个儿童故事,要求如下:

【故事主题】{story_params.get('story_topic', '')}
【主要角色】{story_params.get('main_role', '')}
【故事场景】{story_params.get('scene', '')}
【风格】{story_params.get('style', '温馨有趣')}
【页数要求】必须严格生成{self.target_pages}页内容

【格式要求】
1. 必须严格按照"第X页:"的格式输出
2. 每页2-4句话,适合儿童阅读
3. 故事要有完整的开头、发展、高潮、结尾
4. 不要输出任何其他说明文字

现在请开始创作,直接输出第1页到第{self.target_pages}页的内容:"""
            
            self.progress = 50
            self.status_message = "正在使用AI创作故事..."
            
            story_content = self.llm.generate(
                prompt=prompt,
                system_prompt="你是专业的儿童故事作家。",
                temperature=0.8,
                max_tokens=2000
            )
            
            self.progress = 70
            self.status_message = "正在解析故事内容..."
            
            pages = self._parse_story_pages(story_content)
            
            if not pages:
                raise Exception("故事解析失败")
            
            self.story_pages = pages
            self.progress = 100
            self.current_stage = "completed"
            self.status_message = f"故事生成完成!共{len(pages)}页"
            
            return pages
            
        except Exception as e:
            self.current_stage = "error"
            self.status_message = f"生成失败: {str(e)}"
            raise
        finally:
            self.is_generating = False
    
    def regenerate_page(self, page_index: int, user_feedback: str = None) -> dict:
        """重新生成指定页面"""
        try:
            if page_index < 0 or page_index >= len(self.story_pages):
                return {"success": False, "error": "页面索引无效"}
            
            self.is_generating = True
            self.current_stage = "regenerating"
            self.progress = 50
            self.status_message = f"正在重新生成第 {page_index + 1} 页..."
            
            context_before = self.story_pages[page_index - 1] if page_index > 0 else ""
            context_after = self.story_pages[page_index + 1] if page_index < len(self.story_pages) - 1 else ""
            current_page = self.story_pages[page_index]
            
            feedback_text = f"\n用户反馈: {user_feedback}" if user_feedback else ""
            
            prompt = f"""请重新创作儿童故事的第 {page_index + 1} 页内容。

【上一页内容】
{context_before if context_before else "这是第一页"}

【当前页(需要改进)】
{current_page}

【下一页内容】
{context_after if context_after else "这是最后一页"}
{feedback_text}

【要求】
1. 保持与前后页的连贯性
2. 每页2-4句话,适合儿童阅读
3. 语言生动有趣
4. 只输出新的内容,不要包含"第X页:"这样的标记

请直接输出改进后的内容:"""
            
            new_content = self.llm.generate(
                prompt=prompt,
                system_prompt="你是专业的儿童故事作家。",
                temperature=0.9,
                max_tokens=500
            )
            
            new_content = new_content.strip()
            new_content = re.sub(r'第\s*\d+\s*页\s*[::]?\s*', '', new_content)
            
            self.story_pages[page_index] = new_content
            
            if hasattr(self, 'image_paths') and page_index < len(self.image_paths):
                if self.image_paths[page_index]:
                    old_path = self.image_paths[page_index]
                    if os.path.exists(old_path):
                        try:
                            os.remove(old_path)
                        except:
                            pass
                    self.image_paths[page_index] = None
            
            self.progress = 100
            self.current_stage = "completed"
            self.status_message = f"第 {page_index + 1} 页重新生成完成"
            
            return {
                "success": True,
                "new_content": new_content,
                "page_index": page_index
            }
            
        except Exception as e:
            self.current_stage = "error"
            self.status_message = f"重新生成失败: {str(e)}"
            return {"success": False, "error": str(e)}
        finally:
            self.is_generating = False
    
    def _parse_story_pages(self, story_content: str) -> list:
        pages = []
        pattern = r'第\s*(\d+)\s*页\s*[::]\s*'
        parts = re.split(pattern, story_content)
        
        for i in range(1, len(parts), 2):
            if i + 1 < len(parts):
                content = parts[i + 1].strip()
                content = re.sub(r'\n+', ' ', content)
                content = re.sub(r'\s+', ' ', content)
                
                if content and len(content) > 10:
                    pages.append(content)
        
        if not pages:
            paragraphs = [p.strip() for p in story_content.split('\n') if p.strip()]
            paragraphs = [p for p in paragraphs if len(p) > 20]
            pages = paragraphs[:self.target_pages]
        
        if len(pages) > self.target_pages:
            pages = pages[:self.target_pages]
        
        return pages


class DirectStorySession:
    """直接文本生成故事会话 - 🔥 修复版"""
    def __init__(self, session_id: str, api_key: str):
        self.session_id = session_id
        self.api_key = api_key
        self.story_pages = []
        self.image_paths = []
        self.video_path = None
        self.current_stage = "initial"
        self.is_generating = False
        self.progress = 0
        self.status_message = ""
        self.llm = DashScopeClient(api_key)
        self.uploaded_file_path = None  # 🔥 新增

    def update_progress(self, stage, progress, message):
        self.current_stage = stage
        self.progress = progress
        self.status_message = message
    
    def generate_story_from_text(self, text: str, target_pages: int = 8) -> list:
        """从文本生成故事"""
        try:
            self.is_generating = True
            self.current_stage = "generating"
            self.progress = 30
            
            prompt = f"""请根据下面的描述创作一个儿童故事:

{text}

要求:
1. 生成{target_pages}页的完整故事
2. 每页2-4句话
3. 必须严格按照"第X页:"的格式输出
4. 不要输出任何其他说明文字

现在请开始创作:"""
            
            self.progress = 50
            self.status_message = "正在使用AI创作故事..."
            
            story_content = self.llm.generate(
                prompt=prompt,
                system_prompt="你是专业的儿童故事作家。",
                temperature=0.8,
                max_tokens=2000
            )
            
            self.progress = 70
            pages = self._parse_story_pages(story_content, target_pages)
            
            self.story_pages = pages
            self.progress = 100
            self.current_stage = "completed"
            self.status_message = f"故事生成完成!共{len(pages)}页"
            
            return pages
            
        except Exception as e:
            self.current_stage = "error"
            self.status_message = f"生成失败: {str(e)}"
            raise
        finally:
            self.is_generating = False
    
    def generate_story_from_file(self, file_path: str, target_pages: int = 8) -> list:
        """🔥 新增:从文件生成故事"""
        try:
            self.is_generating = True
            self.current_stage = "generating"
            self.progress = 10
            self.status_message = "正在读取文件..."
            
            # 使用底层的 story_writer
            story_config = config.get('story_writer', {})
            if not story_config:
                raise Exception("配置文件中缺少 story_writer")
            
            from mm_story_agent.base import init_tool_instance
            story_writer = init_tool_instance(story_config)
            
            self.progress = 30
            self.status_message = "正在生成故事..."
            
            # 准备参数
            params = story_config.get('params', {}).copy()
            params['input_type'] = 'long_text'
            params['file_path'] = file_path
            params['target_pages'] = target_pages
            
            # 调用生成
            pages = story_writer.call(params)
            
            self.story_pages = pages
            self.progress = 100
            self.current_stage = "completed"
            self.status_message = f"故事生成完成!共{len(pages)}页"
            
            return pages
            
        except Exception as e:
            self.current_stage = "error"
            self.status_message = f"生成失败: {str(e)}"
            raise
        finally:
            self.is_generating = False
    
    def regenerate_page(self, page_index: int, user_feedback: str = None) -> dict:
        """重新生成指定页面 - 与ConversationalStorySession相同的实现"""
        try:
            if page_index < 0 or page_index >= len(self.story_pages):
                return {"success": False, "error": "页面索引无效"}
            
            self.is_generating = True
            self.current_stage = "regenerating"
            self.progress = 50
            self.status_message = f"正在重新生成第 {page_index + 1} 页..."
            
            context_before = self.story_pages[page_index - 1] if page_index > 0 else ""
            context_after = self.story_pages[page_index + 1] if page_index < len(self.story_pages) - 1 else ""
            current_page = self.story_pages[page_index]
            
            feedback_text = f"\n用户反馈: {user_feedback}" if user_feedback else ""
            
            prompt = f"""请重新创作儿童故事的第 {page_index + 1} 页内容。

【上一页内容】
{context_before if context_before else "这是第一页"}

【当前页(需要改进)】
{current_page}

【下一页内容】
{context_after if context_after else "这是最后一页"}
{feedback_text}

【要求】
1. 保持与前后页的连贯性
2. 每页2-4句话,适合儿童阅读
3. 语言生动有趣
4. 只输出新的内容,不要包含"第X页:"这样的标记

请直接输出改进后的内容:"""
            
            new_content = self.llm.generate(
                prompt=prompt,
                system_prompt="你是专业的儿童故事作家。",
                temperature=0.9,
                max_tokens=500
            )
            
            new_content = new_content.strip()
            new_content = re.sub(r'第\s*\d+\s*页\s*[::]?\s*', '', new_content)
            
            self.story_pages[page_index] = new_content
            
            if hasattr(self, 'image_paths') and page_index < len(self.image_paths):
                if self.image_paths[page_index]:
                    old_path = self.image_paths[page_index]
                    if os.path.exists(old_path):
                        try:
                            os.remove(old_path)
                        except:
                            pass
                    self.image_paths[page_index] = None
            
            self.progress = 100
            self.current_stage = "completed"
            self.status_message = f"第 {page_index + 1} 页重新生成完成"
            
            return {
                "success": True,
                "new_content": new_content,
                "page_index": page_index
            }
            
        except Exception as e:
            self.current_stage = "error"
            self.status_message = f"重新生成失败: {str(e)}"
            return {"success": False, "error": str(e)}
        finally:
            self.is_generating = False
    
    def _parse_story_pages(self, story_content: str, target_pages: int) -> list:
        pages = []
        pattern = r'第\s*(\d+)\s*页\s*[::]\s*'
        parts = re.split(pattern, story_content)
        
        for i in range(1, len(parts), 2):
            if i + 1 < len(parts):
                content = parts[i + 1].strip()
                content = re.sub(r'\n+', ' ', content)
                
                if content and len(content) > 10:
                    pages.append(content)
        
        if len(pages) > target_pages:
            pages = pages[:target_pages]
        
        return pages


def get_api_key():
    """获取API Key"""
    api_key = config.get('image_generation', {}).get('params', {}).get('api_key')
    if not api_key:
        api_key = os.getenv('DASHSCOPE_API_KEY')
    if not api_key:
        api_key = "sk-1d3302438b004fe8ac7ded13d2d95548"
    return api_key


@app.route('/')
def index():
    """主页面"""
    try:
        html_file = 'improved-story-system.html'
        if not os.path.exists(html_file):
            html_file = 'integrated-story-system.html'
        return send_file(html_file, mimetype='text/html')
    except Exception as e:
        return f"Error loading HTML: {str(e)}", 500


@app.route('/health')
def health():
    return jsonify({
        'status': 'healthy',
        'active_sessions': len(user_sessions),
        'api_configured': bool(get_api_key())
    })


@app.route('/api/start-chat-session', methods=['POST'])
def start_chat_session():
    """开始对话式生成会话"""
    try:
        data = request.json
        session_id = str(uuid.uuid4())
        
        api_key = get_api_key()
        if not api_key:
            return jsonify({"success": False, "error": "未配置API Key"}), 500
        
        target_pages = data.get('config', {}).get('target_pages', 8)
        session = ConversationalStorySession(session_id, api_key, target_pages)
        user_sessions[session_id] = session
        
        welcome = f"你好!我是故事创作助手。告诉我,你想创作一个什么样的故事?我们将生成 {target_pages} 页的故事。"
        
        session.dialogue_history.append({
            "role": "assistant",
            "content": welcome,
            "timestamp": datetime.now().isoformat()
        })
        
        save_session(session_id, session)
        
        return jsonify({
            "success": True,
            "session_id": session_id,
            "welcome_message": welcome
        })
    
    except Exception as e:
        print(f"❌ 创建会话失败: {e}")
        return jsonify({"success": False, "error": str(e)}), 500


@app.route('/api/chat', methods=['POST'])
def chat():
    """对话接口"""
    try:
        data = request.json
        session_id = data.get('session_id')
        message = data.get('message')
        target_pages = data.get('target_pages')
        
        if session_id not in user_sessions:
            session = load_session(session_id)
            if session:
                user_sessions[session_id] = session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        session = user_sessions[session_id]
        result = session.chat(message, target_pages)
        save_session(session_id, session)
        
        return jsonify(result)
    
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route('/api/start-direct-session', methods=['POST'])
def start_direct_session():
    """开始直接文本生成会话"""
    try:
        session_id = str(uuid.uuid4())
        api_key = get_api_key()
        
        session = DirectStorySession(session_id, api_key)
        user_sessions[session_id] = session
        save_session(session_id, session)
        
        return jsonify({"success": True, "session_id": session_id})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route('/api/upload-file', methods=['POST'])
def upload_file():
    """🔥 新增:文件上传接口"""
    try:
        if 'file' not in request.files:
            return jsonify({"success": False, "error": "未找到文件"}), 400
        
        file = request.files['file']
        session_id = request.form.get('session_id')
        
        if not session_id:
            return jsonify({"success": False, "error": "缺少session_id"}), 400
        
        if file.filename == '':
            return jsonify({"success": False, "error": "文件名为空"}), 400
        
        if not allowed_file(file.filename):
            return jsonify({
                "success": False, 
                "error": f"不支持的文件类型。支持: {', '.join(ALLOWED_EXTENSIONS)}"
            }), 400
        
        # 保存文件
        filename = secure_filename(file.filename)
        session_folder = UPLOAD_FOLDER / session_id
        session_folder.mkdir(exist_ok=True)
        
        file_path = session_folder / filename
        file.save(str(file_path))
        
        # 保存文件路径到会话
        if session_id in user_sessions:
            session = user_sessions[session_id]
            if isinstance(session, DirectStorySession):
                session.uploaded_file_path = str(file_path)
                save_session(session_id, session)
        
        print(f"✅ 文件上传成功: {file_path}")
        
        return jsonify({
            "success": True,
            "file_path": str(file_path),
            "filename": filename,
            "size": os.path.getsize(file_path)
        })
        
    except Exception as e:
        print(f"❌ 文件上传失败: {e}")
        return jsonify({"success": False, "error": str(e)}), 500


@app.route('/api/regenerate-page', methods=['POST'])
def regenerate_page():
    """重新生成指定页面"""
    try:
        data = request.json
        session_id = data.get('session_id')
        page_index = data.get('page_index')
        user_feedback = data.get('feedback', '')
        
        if session_id not in user_sessions:
            session = load_session(session_id)
            if session:
                user_sessions[session_id] = session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        session = user_sessions[session_id]
        
        def regenerate_in_background():
            try:
                result = session.regenerate_page(page_index, user_feedback)
                save_session(session_id, session)
            except Exception as e:
                print(f"❌ 重新生成失败: {e}")
                session.current_stage = "error"
                session.status_message = f"重新生成失败: {str(e)}"
                save_session(session_id, session)
        
        thread = threading.Thread(target=regenerate_in_background)
        thread.daemon = True
        thread.start()
        
        return jsonify({"success": True, "session_id": session_id})
        
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route('/api/generate-from-chat', methods=['POST'])
def generate_from_chat():
    """从对话生成故事"""
    try:
        data = request.json
        session_id = data.get('session_id')
        story_params = data.get('story_params')
        
        if session_id not in user_sessions:
            session = load_session(session_id)
            if session:
                user_sessions[session_id] = session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        session = user_sessions[session_id]
        
        def generate_in_background():
            try:
                pages = session.generate_story(story_params)
                save_session(session_id, session)
            except Exception as e:
                print(f"❌ 生成失败: {e}")
                save_session(session_id, session)
        
        thread = threading.Thread(target=generate_in_background)
        thread.daemon = True
        thread.start()
        
        return jsonify({"success": True, "session_id": session_id})
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500

@app.route('/api/generate-from-text', methods=['POST'])
def generate_from_text():
    """🔥 修复版:从长文本生成故事"""
    try:
        data = request.json
        session_id = data.get('session_id')
        text = data.get('text', '').strip()
        target_pages = data.get('target_pages', 8)
        use_uploaded_file = data.get('use_uploaded_file', False)
        
        # ✅ 会话检查和获取（完整版）
        if session_id not in user_sessions:
            loaded_session = load_session(session_id)
            if loaded_session:
                user_sessions[session_id] = loaded_session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        # 🔥 关键：必须在这里定义 session 变量
        session = user_sessions[session_id]
        
        # 验证会话类型
        if not isinstance(session, DirectStorySession):
            return jsonify({"success": False, "error": "会话类型错误，请使用文本模式"}), 400
        
        # 🔥 关键修复1: 获取 session 对象
        if session_id not in user_sessions:
            session = load_session(session_id)
            if session:
                user_sessions[session_id] = session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        # 🔥 关键修复2: 将 session 赋值给局部变量
        session = user_sessions[session_id]
        
        # 验证 session 类型
        if not isinstance(session, DirectStorySession):
            return jsonify({"success": False, "error": "会话类型错误"}), 400
        
        print(f"📝 从文本生成故事: hasText={bool(text)}, hasFile={bool(use_uploaded_file)}, sessionId={session_id}")
        
        def generate_in_background():
            try:
                # 现在可以安全地使用 session 了
                if use_uploaded_file and hasattr(session, 'uploaded_file_path') and session.uploaded_file_path:
                    print(f"📄 使用上传的文件: {session.uploaded_file_path}")
                    pages = session.generate_story_from_file(
                        session.uploaded_file_path, 
                        target_pages
                    )
                elif text:
                    print(f"📝 使用文本输入 ({len(text)} 字符)")
                    pages = session.generate_story_from_text(text, target_pages)
                else:
                    raise Exception("未提供文本或文件")
                
                save_session(session_id, session)
           
            except Exception as e:
                print(f"❌ 生成失败: {e}")
                import traceback
                traceback.print_exc()
                # 现在 session 已经定义了，不会报错
                session.current_stage = "error"
                session.status_message = f"生成失败: {str(e)}"
                save_session(session_id, session)
        
        thread = threading.Thread(target=generate_in_background)
        thread.daemon = True
        thread.start()
        
        return jsonify({"success": True, "session_id": session_id})
        
    except Exception as e:
        print(f"❌ 请求处理失败: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({"success": False, "error": str(e)}), 500
    
@app.route('/api/chat-session-status/<session_id>', methods=['GET'])
def chat_session_status(session_id):
    """获取会话状态"""
    if session_id not in user_sessions:
        session = load_session(session_id)
        if session:
            user_sessions[session_id] = session
        else:
            return jsonify({"success": False, "error": "会话不存在"}), 404
    
    session = user_sessions[session_id]
    
    generated_images = []
    if hasattr(session, 'image_paths') and session.image_paths:
        for idx, img_path in enumerate(session.image_paths):
            if img_path and os.path.exists(img_path):
                generated_images.append({
                    "index": idx,
                    "url": f"/api/image/{session_id}/{idx}"
                })
    
    return jsonify({
        "success": True,
        "stage": session.current_stage,
        "progress": session.progress,
        "status_message": session.status_message,
        "is_generating": session.is_generating,
        "has_story": len(session.story_pages) > 0,
        "story_pages": session.story_pages if not session.is_generating else [],
        "generated_images": generated_images,
        "has_video": bool(getattr(session, 'video_path', None))
    })


@app.route('/api/image/<session_id>/<int:image_index>')
def get_image(session_id, image_index):
    """获取生成的图片"""
    if session_id not in user_sessions:
        session = load_session(session_id)
        if session:
            user_sessions[session_id] = session
        else:
            return jsonify({"error": "会话不存在"}), 404
    
    session = user_sessions[session_id]
    
    if not hasattr(session, 'image_paths') or image_index >= len(session.image_paths):
        return jsonify({"error": "图片不存在"}), 404
    
    img_path = session.image_paths[image_index]
    if not img_path or not os.path.exists(img_path):
        return jsonify({"error": "图片文件不存在"}), 404
    
    return send_file(img_path, mimetype='image/png')


@app.route('/api/video/<session_id>')
def get_video(session_id):
    """获取生成的视频"""
    if session_id not in user_sessions:
        session = load_session(session_id)
        if session:
            user_sessions[session_id] = session
        else:
            return jsonify({"error": "会话不存在"}), 404
    
    session = user_sessions[session_id]
    
    if not hasattr(session, 'video_path') or not session.video_path:
        story_dir = Path(config.get('story_dir', 'generated_stories/example'))
        video_path = story_dir / "output.mp4"
        if video_path.exists():
            return send_file(str(video_path), mimetype='video/mp4')
        return jsonify({"error": "视频未生成"}), 404
    
    video_path = session.video_path
    if not os.path.exists(video_path):
        return jsonify({"error": "视频文件不存在"}), 404
    
    return send_file(video_path, mimetype='video/mp4')


@app.route('/api/generate-images', methods=['POST'])
def generate_images():
    """生成故事配图"""
    try:
        data = request.json
        session_id = data.get('session_id')
        
        if session_id not in user_sessions:
            session = load_session(session_id)
            if session:
                user_sessions[session_id] = session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        session = user_sessions[session_id]
        
        if not session.story_pages:
            return jsonify({"success": False, "error": "请先生成故事"}), 400
        
        def generate_images_background():
            try:
                session.update_progress("generating_images", 10, "正在生成图片...")
                
                from mm_story_agent.base import init_tool_instance
                image_config = config.get('image_generation', {})
                
                if 'params' not in image_config:
                    image_config['params'] = {}
                image_config['params']['api_key'] = get_api_key()
                
                image_agent = init_tool_instance(image_config)
                
                session.update_progress("generating_images", 30, "正在分析故事...")
                role_dict = image_agent.extract_role_from_story(session.story_pages)
                
                session.update_progress("generating_images", 40, "正在生成提示词...")
                prompts = image_agent.generate_image_prompt_from_story(session.story_pages)
                
                story_dir = Path(config.get('story_dir', 'generated_stories/example'))
                image_dir = story_dir / "image"
                image_dir.mkdir(parents=True, exist_ok=True)
                
                session.update_progress("generating_images", 50, "使用API生成图片中...")
                
                params = {
                    'pages': session.story_pages,
                    'save_path': image_dir,
                    'style_name': image_config.get('params', {}).get('style_name', 'Storybook')
                }
                
                result = image_agent.call(params)
                images = result.get('generation_results', [])
                
                image_paths = [img for img in images if img and os.path.exists(img)]
                session.image_paths = image_paths
                
                session.update_progress("completed", 100, 
                    f"图片生成完成!成功 {len(image_paths)}/{len(session.story_pages)} 张")
                
                save_session(session.session_id, session)
                
            except Exception as e:
                print(f"❌ 图片生成失败: {e}")
                import traceback
                traceback.print_exc()
                session.update_progress("error", 0, f"图片生成失败: {str(e)}")
                save_session(session.session_id, session)
        
        thread = threading.Thread(target=generate_images_background)
        thread.daemon = True
        thread.start()
        
        return jsonify({"success": True})
        
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route('/api/compose-video', methods=['POST'])
def compose_video():
    """合成视频"""
    try:
        data = request.json
        session_id = data.get('session_id')
        
        if session_id not in user_sessions:
            session = load_session(session_id)
            if session:
                user_sessions[session_id] = session
            else:
                return jsonify({"success": False, "error": "会话不存在"}), 404
        
        session = user_sessions[session_id]
        
        if not session.story_pages:
            return jsonify({"success": False, "error": "请先生成故事"}), 400
        
        def compose_video_background():
            try:
                session.update_progress("composing_video", 10, "正在准备视频合成...")
                
                story_dir = Path(config.get('story_dir', 'generated_stories/example'))
                
                session.update_progress("composing_video", 50, "正在合成视频...")
                
                video_config = config.get('video_compose', {})
                from mm_story_agent.base import init_tool_instance
                video_agent = init_tool_instance(video_config)
                
                params = video_config.get('params', {}).copy()
                params['pages'] = session.story_pages
                params['story_dir'] = str(story_dir)
                
                if session.image_paths:
                    params['image_paths'] = [str(p) for p in session.image_paths if p]
                
                result = video_agent.call(params)
                
                output_path = story_dir / "output.mp4"
                
                import time
                max_wait = 10
                waited = 0
                while not output_path.exists() and waited < max_wait:
                    time.sleep(1)
                    waited += 1
                
                if output_path.exists():
                    session.video_path = str(output_path)
                    session.update_progress("completed", 100, f"视频已保存: {output_path}")
                else:
                    raise Exception("视频文件未生成")
                
                save_session(session.session_id, session)
                
            except Exception as e:
                print(f"❌ 视频合成失败: {e}")
                import traceback
                traceback.print_exc()
                session.update_progress("error", 0, f"视频合成失败: {str(e)}")
                save_session(session.session_id, session)
        
        thread = threading.Thread(target=compose_video_background)
        thread.daemon = True
        thread.start()
        
        return jsonify({"success": True})
        
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


if __name__ == '__main__':
    print("🚀 启动改进版故事生成系统...")
    print("📁 工作目录:", os.getcwd())
    print("🌐 访问地址: http://127.0.0.1:5000")
    print("🔧 健康检查: http://127.0.0.1:5000/health")
    
    api_key = get_api_key()
    if api_key:
        print(f"✅ API Key已配置: {api_key[:10]}...")
    else:
        print("⚠️ 警告: API Key未配置")
    
    print("\n✨ 新功能:")
    print("  1. ✅ 可自定义页面数量(4-20页)")
    print("  2. ✅ 视频在线预览和下载")
    print("  3. ✅ 网页文件上传功能 (新修复)")
    print("  4. ✅ 对话和文本两种模式")
    print("  5. ✅ 会话自动持久化")
    print("  6. ✅ 实时进度显示")
    print("  7. ✨ 局部页面重新生成")
    
    app.run(debug=True, port=5000)